/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2024-2024.
 * Description: RTOS Process Context asm header
 * Create: 2024-07-13
 */
#include <linux/sched.h>

#include <linux/bug.h>
#include <linux/compiler.h>
#include <linux/atomic.h>
#include <linux/jump_label.h>
#include <linux/rtos_rtpc.h>
#include <linux/rtos_rtpc_rpc.h>

#include <asm/rtos_rtpc.h>
#include <asm/rtos_rtpc_rpc.h>

#include "../kernel/audit.h"

#include "rtpc.h"

__attribute__((cold)) noinline long ret_call_fast_need_exit_handler(void);
__attribute__((cold)) noinline void do_fixup_need_exit_handler(void);

/*
 * We use __always_inline here will not increase the size
 * because this function will only be used once
 */
static __always_inline
void do_call_fast(struct task_struct *const client_task,
		struct task_struct *const server_task,
		struct rpc_service_info *const service_info)
{
	/* help the compiler optimizes the code */
	if (unlikely(client_task == server_task))
		__builtin_unreachable();

	WRITE_ONCE(service_info->calling_client_task, client_task);

	rtpc_context_switch_user(client_task, server_task);

	if (static_branch_unlikely(&sysctl_rpc_kernel_stats))
		rtpc_stat_call();
}

/*
 * We use __always_inline here will not increase the size
 * because this function will only be used once
 */
static __always_inline
long ret_call_fast(struct task_struct *const client_task,
		struct task_struct *const server_task,
		struct rpc_service_info *const service_info)
{
	/* help the compiler optimizes the code */
	if (unlikely(client_task == server_task))
		__builtin_unreachable();

	WRITE_ONCE(service_info->calling_client_task, NULL);

	rtpc_set_current(server_task);
	rtpc_context_switch_user(server_task, client_task);
	rtpc_set_current(client_task);

	if (unlikely(service_info->need_exit))
		return ret_call_fast_need_exit_handler();
	else
		return 0;
}

/*
 * We use __always_inline here will not increase the size
 * because this function will only be used once
 */
static __always_inline
void do_fixup(struct task_struct *const client_task,
		struct task_struct *const server_task,
		struct rpc_service_info *const service_info, const size_t reason)

{
	if (is_pgtable_diff_check(server_task)) {
		rpc_emerg("unexpected pgdiff check in fixup!");
		BUG();
	}

	/* help the compiler optimizes the code */
	if (unlikely(client_task == server_task))
		__builtin_unreachable();

	rtpc_spin_lock(&service_info->lock);

#ifdef CONFIG_RTOS_RTPC_RPC_DEBUG_BUILD
	if (service_info->need_exit & (unsigned int)EXIT_FIXUP)
		BUG_ON(!service_info->server_need_wake || !service_info->calling_client_need_wake);

	BUG_ON(client_task != service_info->calling_client_task);
	BUG_ON(server_task != service_info->server_task);
	BUG_ON(service_info->fixup_done);

	/* need_exit normal flag should not be set */
	BUG_ON(service_info->need_exit & (~((unsigned int)KILL_FIXUP | (unsigned int)EXIT_FIXUP)));

	BUG_ON(server_task->audit_context && server_task->audit_context->in_syscall);
	BUG_ON(client_task->audit_context && client_task->audit_context->in_syscall);
#endif
	rtpc_sched_switch(client_task, server_task);
	rtpc_context_switch_kernel(client_task, server_task);

	rtpc_set_current(server_task);

	rpc_create_context(client_task);

	WRITE_ONCE(service_info->fixup_done, true);
	if (unlikely(service_info->need_exit))
		do_fixup_need_exit_handler();

	rtpc_spin_unlock(&service_info->lock);

	if (static_branch_unlikely(&sysctl_rpc_kernel_stats))
		rtpc_stat_fixup(reason);
}

/*
 * We use __always_inline here will not increase the size
 * because this function will only be used once
 */
static __always_inline
void ret_call_exception(struct task_struct *const client_task,
		struct task_struct *const server_task,
		struct rpc_service_info *const service_info)
{
	/* help the compiler optimizes the code */
	if (unlikely(client_task == server_task))
		__builtin_unreachable();

	WRITE_ONCE(service_info->fixup_done, false);
	WRITE_ONCE(service_info->calling_client_task, NULL);

#ifdef CONFIG_RTOS_RTPC_RPC_DEBUG_BUILD
	BUG_ON(client_task->audit_context && client_task->audit_context->in_syscall);
	BUG_ON(server_task->audit_context && server_task->audit_context->in_syscall);
#endif
	rtpc_sched_switch(server_task, client_task);
	rtpc_context_switch_user(server_task, client_task);
	rtpc_context_switch_kernel(server_task, client_task);

	rtpc_set_current(client_task);

	rpc_create_context(server_task);
}
